# routers.py
from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi.responses import StreamingResponse
from typing import Dict, Any, Optional
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.store.postgres.aio import AsyncPostgresStore
from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata
from depends.dependencies import get_postgres_saver, get_postgres_store
from agent.agent import LangGraphChatbot
from schema.chat import ChatRequestBody
import asyncio
import json
from typing import AsyncIterator

chat_router = APIRouter()

@chat_router.post("/chat", summary="聊天")
async def chat(
    chat_body: ChatRequestBody,
    request: Request,
    saver: AsyncPostgresSaver = Depends(get_postgres_saver)
):
    """保存LangGraph检查点数据"""
    try:
        # 从请求体获取JSON数据
        messages = chat_body.message
        thread_id = chat_body.thread_id

       
        # 创建自定义 agent
        agent = LangGraphChatbot(checkpointer=saver)

        config = {"configurable": {"thread_id": f"{thread_id}"}}

        data = {"res": ""}
        
        async for event in agent.graph.astream(
            {"messages": [("user", f"{messages}")]}, config):
            print(event)
            # 获取最终响应
            data["res"] = event["chatbot"]["messages"][-1].content
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid request data: {e}")
    
    return data

async def stream_generator(chatbot: LangGraphChatbot, message: str, thread_id: str) -> AsyncIterator[str]:
    """异步流式生成器
    
    生成 SSE (Server-Sent Events) 格式的数据流
    """
    try:
        config = {"configurable": {"thread_id": f"{thread_id}"}}
        
        # 使用 astream_events 获取流式事件
        async for event in chatbot.graph.astream_events(
            {"messages": [("user", message)]}, 
            config,
            version="v2"
        ):
            # 处理不同类型的事件
            kind = event.get("event")
            
            # 处理 LLM 流式输出
            if kind == "on_chat_model_stream":
                content = event["data"]["chunk"].content
                if content:
                    data = {
                        "content": content,
                        "thread_id": thread_id,
                        "type": "message"
                    }
                    yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
            
            # 处理工具调用
            elif kind == "on_tool_start":
                tool_name = event.get("name", "unknown")
                data = {
                    "content": f"{tool_name}",
                    "thread_id": thread_id,
                    "type": "tool"
                }
                yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
            
            # 短暂延迟
            await asyncio.sleep(0.001)
        
        # 发送结束标记
        yield f"data: {json.dumps({'type': 'done'})}\n\n"
        
    except Exception as e:
        # 发送错误信息
        error_data = {
            "type": "error",
            "content": str(e)
        }
        yield f"data: {json.dumps(error_data)}\n\n"

        

@chat_router.post("/chat/stream")
async def chat_stream(
    chat_body: ChatRequestBody, 
    saver: AsyncPostgresSaver = Depends(get_postgres_saver)
):
    """流式聊天接口
    
    Args:
        chat_body: 聊天请求体
        
    Returns:
        SSE 流式响应
    """
    try:
        # 创建自定义 agent
        agent = LangGraphChatbot(checkpointer=saver)

        # 从请求体获取JSON数据
        messages = chat_body.message
        thread_id = chat_body.thread_id
        
        return StreamingResponse(
            stream_generator(agent, messages, thread_id),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no"  # 禁用 Nginx 缓冲
            }
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))



@chat_router.get("/history/{thread_id}")
async def get_history(
    thread_id: str,
    saver: AsyncPostgresSaver = Depends(get_postgres_saver)):
    """获取对话历史
    
    Args:
        thread_id: 会话ID
        
    Returns:
        消息历史列表
    """
    # 创建自定义 agent
    agent = LangGraphChatbot(checkpointer=saver)
    try:
        loop = asyncio.get_event_loop()
        history = await loop.run_in_executor(
            None,
            agent.get_history,
            thread_id
        )
        
        # 格式化历史消息
        formatted_history = []
        for msg in history:
            formatted_history.append({
                "role": msg.type if hasattr(msg, 'type') else "unknown",
                "content": msg.content if hasattr(msg, 'content') else str(msg)
            })
        
        return {"thread_id": thread_id, "messages": formatted_history}
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@chat_router.delete("/history/{thread_id}")
async def clear_history(
    thread_id: str,
    saver: AsyncPostgresSaver = Depends(get_postgres_saver)):
    """清除对话历史
    
    Args:
        thread_id: 会话ID
    """
    # 创建自定义 agent
    agent = LangGraphChatbot(checkpointer=saver)
    try:
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(
            None,
            agent.clear_history,
            thread_id
        )
        
        return {"message": "历史记录已清除", "thread_id": thread_id}
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))





    